import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl


def srd_euler(M, x0, dt, theta, i, k, sigma):
    """
    扩散模型：偏移量 + 扩散
    x(t)： 连续：卡方分布；离散：正态分布
    dx(t) = k * (theta - x(t)) * dt + sigma * sqrt(x(t)) * w(t)

    解方程：欧拉数值方法迭代近似计算
    x(t) = x(s) + k * (theta - max(x(s), 0)) * detat + sigmal * sqrt(max(x(s), 0) * detat) * w(t)
    x(t) = max(x(t), 0)
    :param M:
    :param x0:
    :param dt:
    :param theta:
    :param i:
    :param k:
    :param sigma:
    :return:
    """
    xh = np.zeros((M + 1, i))
    x1 = np.zeros_like(xh)
    xh[0], x1[0] = x0, x0
    for t in range(1, M + 1):
        xh[t] = xh[t - 1] + \
                k * (theta - np.maximum(xh[t - 1], 0)) * dt + \
                sigma * np.sqrt(np.maximum(xh[t - 1], 0) * dt) * np.random.standard_normal(i)

    x1 = np.maximum(xh, 0)

    return x1


def diffusion_main():
    """
    平方根扩展模拟
    :return:
    """
    # matplotlib中文显示方块
    mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
    mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

    s0 = 100
    r = 0.05
    sigma = 0.25
    T = 2.0
    x0 = 0
    k = 1.8
    theta = 0.24
    M = 50
    dt = T / M
    i = 100000

    # 计算
    x1 = srd_euler(M, x0, dt, theta, i, k, sigma)

    # 画图
    # plt.hist(x1[-1], bins=30, alpha=0.85)
    # plt.xlabel('数值')
    # plt.ylabel('频次')
    # plt.grid(False)

    plt.plot(x1[:, :5])
    plt.title('平方根扩展模拟')
    plt.xlabel('时间')
    plt.ylabel('指数level')
    plt.grid(True)

    plt.show()


if __name__ == '__main__':
    diffusion_main()
